# coding:utf-8
import sys
import os
import re
import maya.OpenMaya as om
import maya.OpenMayaMPx as ompx

kPluginNodeTypeName = "filePathResolver"

kFilePathResolverId = om.MTypeId(0x88000)


# 节点定义
class FilePathResolver(ompx.MPxNode):
    # class variables
    filePath = om.MObject()
    outFilePath = om.MObject()

    def __init__(self):
        ompx.MPxNode.__init__(self)

    def compute(self, plug, dataBlock):
        if plug == FilePathResolver.outFilePath:
            dataHandle = dataBlock.inputValue(FilePathResolver.filePath)
            inputFilePath = dataHandle.asString()
            # 开始解析该路径
            outputFilePath = self.parserFilePath(inputFilePath)
            # 返回值
            outputHandle = dataBlock.outputValue(FilePathResolver.outFilePath)
            outputHandle.setString(outputFilePath)
            # 标记plug已经不需要计算了
            dataBlock.setClean(plug)
        return om.kUnknownParameter

    # creator
    @staticmethod
    def nodeCreator():
        return ompx.asMPxPtr(FilePathResolver())

    # initializer
    @staticmethod
    def nodeInitializer():
        data = om.MFnStringData()
        # input
        nAttr = om.MFnTypedAttribute()
        FilePathResolver.filePath = nAttr.create("filePath", "fp", om.MFnData.kString, data.create(""), )
        nAttr.setStorable(1)
        nAttr.setUsedAsFilename(1)
        nAttr.setConnectable(0)
        # output
        nAttr = om.MFnTypedAttribute()
        FilePathResolver.outFilePath = nAttr.create("outFilePath", "ofp", om.MFnData.kString, data.create(""), )
        nAttr.setHidden(1)
        nAttr.setStorable(1)
        nAttr.setWritable(0)
        # 影响
        FilePathResolver.addAttribute(FilePathResolver.filePath)
        FilePathResolver.addAttribute(FilePathResolver.outFilePath)
        FilePathResolver.attributeAffects(FilePathResolver.filePath, FilePathResolver.outFilePath)

    @staticmethod
    def parserFilePath(inputPath):
        # 检查该文件路径
        envPath, envs = FilePathResolver.__parse_arnold_syntex_env_file_path(inputPath)
        # 解析该路径
        outFilePath = om.MFileObject.getResolvedFullName(envPath)
        return outFilePath

    @staticmethod
    def __parse_arnold_syntex_env_file_path(file_path):
        # 解释arnold环境变量
        # 解析路径中的[env]环境变量
        new_path = os.path.normpath(file_path).replace(os.path.sep, os.path.altsep)
        found_env = []
        # 找到arnoldEnv
        arnold_envs = re.findall("\[(\w+)\]", new_path)
        if arnold_envs:
            arnold_envs_mapping = {}
            for e in list(set(arnold_envs)):
                if e not in os.environ:
                    return None, arnold_envs
                else:
                    # 收集信息
                    found_env.append(e)
                    #
                    arnold_envs_mapping[e] = os.path.normpath(os.environ[e]).replace(os.path.sep, os.path.altsep)
            # 更新path
            for source, replaced in sorted(arnold_envs_mapping.items(), key=lambda x: x[0]):
                new_path = new_path.replace("[{}]".format(source), replaced)
        return new_path, list(set(found_env))


# initialize the script plug-in
def initializePlugin(mobject):
    mplugin = ompx.MFnPlugin(mobject)
    try:
        mplugin.registerNode(kPluginNodeTypeName, kFilePathResolverId, FilePathResolver.nodeCreator,
                             FilePathResolver.nodeInitializer)
    except:
        sys.stderr.write("Failed to register node: %s" % kPluginNodeTypeName)
        raise


# uninitialize the script plug-in
def uninitializePlugin(mobject):
    mplugin = ompx.MFnPlugin(mobject)
    try:
        mplugin.deregisterNode(kFilePathResolverId)
    except:
        sys.stderr.write("Failed to deregister node: %s" % kPluginNodeTypeName)
        raise
